"""WebSocket client for voice_assistant_contro."""
from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable

import aiohttp
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, Event
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import (
    API_WEBSOCKET_PATH,
    SERVER_WS_HOST,
    WEBSOCKET_HEARTBEAT_INTERVAL,
    WEBSOCKET_RECONNECT_INTERVAL,
    SERVER_HOST,
    API_LOGIN_PATH,
)

_LOGGER = logging.getLogger(__name__)

class VoiceAssistantWebSocket:
    """WebSocket client for voice_assistant_contro."""

    def __init__(
        self, 
        hass: HomeAssistant, 
        entry: ConfigEntry,
        message_handler: Callable[[dict[str, Any]], None],
        connection_state_callback: Callable[[bool, bool], None] | None = None,
    ) -> None:
        """初始化 WebSocket 客户端."""
        self.hass = hass
        self.entry = entry
        self.session = async_get_clientsession(hass)
        self.ws_client: aiohttp.ClientWebSocketResponse | None = None
        self.heartbeat_task: asyncio.Task | None = None
        self.reconnect_task: asyncio.Task | None = None
        self.message_handler_task: asyncio.Task | None = None
        self.closing = False
        self._connection_lock = asyncio.Lock()
        self._connected = False
        self._connecting = False
        self._message_handler = message_handler
        self._connection_state_callback = connection_state_callback
        
        # 获取配置
        self.username = entry.data["username"]
        self.password = entry.data["password"]

    def _update_connection_state(self) -> None:
        """更新连接状态."""
        if self._connection_state_callback:
            self._connection_state_callback(self._connected, self._connecting)

    async def _start_reconnect(self) -> None:
        """如果未运行则启动重连任务."""
        if not self.closing and not self.reconnect_task:
            self.reconnect_task = asyncio.create_task(self._reconnect())
        
    async def connect(self) -> None:
        """连接到 WebSocket 服务器."""
        async with self._connection_lock:
            if self._connected or self._connecting:
                return
                
            self._connecting = True
            self._update_connection_state()
                
            # 确保之前的连接已经完全关闭
            await self._cleanup_connection(start_reconnect=False)
            
            try:    
                # 建立 WebSocket 连接
                ws_url = f"{SERVER_WS_HOST}{API_WEBSOCKET_PATH}?username={self.username}&password={self.password}"
                self.ws_client = await self.session.ws_connect(ws_url)
                self._connected = True
                self._connecting = False
                self._update_connection_state()
                
                # 发送初始化消息
                await self.send_message({
                    "type": "init",
                    "data": {
                        "client_id": self.entry.entry_id
                    }
                })
                
                # 启动心跳任务
                if not self.heartbeat_task:
                    self.heartbeat_task = asyncio.create_task(self._heartbeat())
                
                # 启动消息处理任务
                if not self.message_handler_task:
                    self.message_handler_task = asyncio.create_task(self._handle_messages())
                
                _LOGGER.info("WebSocket 连接成功")
                
            except Exception as err:
                _LOGGER.error("连接到 WebSocket 服务器失败: %s", err)
                await self._cleanup_connection(start_reconnect=True)

    async def _cleanup_connection(self, start_reconnect: bool = False) -> None:
        """清理当前连接和所有相关任务."""
        self._connected = False
        self._connecting = False
        self._update_connection_state()
        
        # 取消心跳任务
        if self.heartbeat_task:
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass
            self.heartbeat_task = None
            
        # 取消消息处理任务
        if self.message_handler_task:
            self.message_handler_task.cancel()
            try:
                await self.message_handler_task
            except asyncio.CancelledError:
                pass
            self.message_handler_task = None
            
        # 关闭 WebSocket 连接
        if self.ws_client:
            if not self.ws_client.closed:
                await self.ws_client.close()
            self.ws_client = None

        # 根据参数决定是否启动重连
        if start_reconnect:
            await self._start_reconnect()
                
    async def disconnect(self, event: Event | None = None) -> None:
        """断开与 WebSocket 服务器的连接."""
        self.closing = True
        
        # 取消重连任务
        if self.reconnect_task:
            self.reconnect_task.cancel()
            try:
                await self.reconnect_task
            except asyncio.CancelledError:
                pass
            self.reconnect_task = None
            
        # 清理连接
        await self._cleanup_connection(start_reconnect=False)
            
    async def _reconnect(self) -> None:
        """重新连接到 WebSocket 服务器."""
        try:
            while not self.closing:
                _LOGGER.info("将在 %d 秒后尝试重新连接", WEBSOCKET_RECONNECT_INTERVAL)
                await asyncio.sleep(WEBSOCKET_RECONNECT_INTERVAL)
                if not self.closing:
                    await self.connect()
                    if self._connected:
                        break
        finally:
            self.reconnect_task = None
            
    async def _heartbeat(self) -> None:
        """发送定期心跳消息."""
        try:
            while True:
                if not self._connected or not self.ws_client or self.ws_client.closed:
                    break
                    
                try:
                    await self.send_message({
                        "type": "ping",
                        "timestamp": int(self.hass.loop.time())
                    })
                    await asyncio.sleep(WEBSOCKET_HEARTBEAT_INTERVAL)
                except Exception as err:
                    _LOGGER.error("心跳错误: %s", err)
                    break
        finally:
            # 如果心跳失败，触发清理和重连
            if not self.closing and self._connected:
                await self._cleanup_connection(start_reconnect=True)
                
    async def _handle_messages(self) -> None:
        """处理传入的 WebSocket 消息."""
        if not self.ws_client:
            return
            
        try:
            async for msg in self.ws_client:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._message_handler(data)
                    except json.JSONDecodeError as err:
                        _LOGGER.error("解码消息失败: %s", err)
                elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break
        except Exception as err:
            _LOGGER.error("WebSocket 消息处理错误: %s", err)
        finally:
            if not self.closing:
                await self._cleanup_connection(start_reconnect=True)

    async def send_message(self, message: dict[str, Any]) -> None:
        """发送消息到服务器."""
        if not self.ws_client or self.ws_client.closed:
            return
            
        try:
            await self.ws_client.send_json(message)
        except Exception as err:
            _LOGGER.error("发送消息失败: %s", err) 